package com.example.shardingjdbc.config;

import com.google.common.collect.Range;
import org.apache.shardingsphere.api.sharding.complex.ComplexKeysShardingAlgorithm;
import org.apache.shardingsphere.api.sharding.complex.ComplexKeysShardingValue;

import java.text.SimpleDateFormat;
import java.util.*;

public class MonthBasedDatabaseComplexShardingAlgorithm implements ComplexKeysShardingAlgorithm<Date> {

    private static final String SHARDING_COLUMN = "data_time";

    @Override
    public Collection<String> doSharding(Collection<String> availableTargetNames, ComplexKeysShardingValue<Date> shardingValue) {
        Set<String> result = new HashSet<>();
        System.out.println("1111111111111111111111111111");
        // 精确值匹配
        if (shardingValue.getColumnNameAndShardingValuesMap() != null && !shardingValue.getColumnNameAndShardingValuesMap().isEmpty()) {
            Collection<Date> values = getShardingValues(shardingValue.getColumnNameAndShardingValuesMap(), SHARDING_COLUMN);
            for (Date date : values) {
                int month = getMonthFromDate(date);
                String target = calculateDatabase(month, availableTargetNames);
                result.add(target);
            }
        }

        // 范围匹配
        else if (shardingValue.getColumnNameAndRangeValuesMap() != null && !shardingValue.getColumnNameAndRangeValuesMap().isEmpty()) {
            Range<Date> range = (Range<Date>) shardingValue.getColumnNameAndRangeValuesMap().get(SHARDING_COLUMN);
            List<Date> datesInRange = getDatesInRange(range);

            for (Date date : datesInRange) {
                int month = getMonthFromDate(date);
                String target = calculateDatabase(month, availableTargetNames);
                result.add(target);
            }
        }

        return result;
    }

    /**
     * 获取指定列的所有值集合
     */
    private Collection<Date> getShardingValues(Map<String, Collection<Date>> columnMap, String key) {
        Collection<Date> result = new ArrayList<>();
        if (columnMap.containsKey(key)) {
            result.addAll(columnMap.get(key));
        } else if (columnMap.containsKey(key.toLowerCase())) {
            result.addAll(columnMap.get(key.toLowerCase()));
        } else if (columnMap.containsKey(key.toUpperCase())) {
            result.addAll(columnMap.get(key.toUpperCase()));
        }
        return result;
    }

    /**
     * 根据月份计算应路由到哪个数据库（ds0/ds1）
     */
    private String calculateDatabase(int month, Collection<String> availableTargetNames) {
        int dbIndex = 1 - (month % 2); // 奇数月在 ds0，偶数月在 ds1
        String target = "ds" + dbIndex;

        if (!availableTargetNames.contains(target)) {
            throw new IllegalArgumentException("No database found for month: " + month + ", target: " + target);
        }

        return target;
    }

    /**
     * 获取月份
     */
    private int getMonthFromDate(Date date) {
        Calendar calendar = Calendar.getInstance();
        calendar.setTime(date);
        return calendar.get(Calendar.MONTH) + 1; // 月份从0开始
    }

    /**
     * 模拟按月拆分时间范围
     */
    private List<Date> getDatesInRange(Range<Date> range) {
        List<Date> result = new ArrayList<>();
        Calendar start = Calendar.getInstance();
        start.setTime(range.lowerEndpoint());
        Calendar end = Calendar.getInstance();
        end.setTime(range.upperEndpoint());

        while (!start.after(end)) {
            result.add(start.getTime());
            start.add(Calendar.MONTH, 1);
        }

        return result;
    }
}
